import os
import chromadb
import ollama
import jieba
from dotenv import load_dotenv

load_dotenv()

class PlanningRAG:
    """
    A RAG class that uses Ollama for embeddings and ChromaDB for storage and retrieval.
    It provides technical knowledge for the planning agent.
    """
    def __init__(self, collection_name="planning_rag_collection", embedding_model="bge-m3:latest"):
        """
        Initializes the RAG system.

        Args:
            collection_name (str): The name of the ChromaDB collection.
            embedding_model (str): The name of the Ollama model for embeddings.
        """
        self.embedding_model = embedding_model
        self.cut_for_search = True
        self.cut_words = True
        self.cut_all = True

        # Initialize ChromaDB client
        self.client = chromadb.Client()
        
        # Create or get the collection
        if collection_name in [c.name for c in self.client.list_collections()]:
            self.collection = self.client.get_collection(name=collection_name)
        else:
            self.collection = self.client.create_collection(name=collection_name)
            self._populate_initial_knowledge()

    def get_cut_words(self, doc: str, debug=False):
        if self.cut_words:
            if self.cut_for_search:
                cut_func = "cut_for_search"
                doc = " ".join(jieba.cut_for_search(doc))
            else:
                cut_func = "cut"
                if self.cut_all:
                    cut_func += "_all"
                else:
                    cut_func += "_1"
                doc = " ".join(jieba.cut(doc, cut_all=self.cut_all))
            if debug:
                print(f"~~~~~ get_cut_words --- cut_func[{cut_func}] --- doc:", doc)
        return doc

    def _populate_initial_knowledge(self):
        """
        Populates the vector database with some initial software development knowledge.
        """
        _documents = [
            "北京是中国的首都，是一个具有悠久历史文化的城市。",
            "上海是中国最大的经济中心，是一个国际化大都市。",
            "深圳被称为中国的硅谷，是改革开放的窗口。",
            "杭州以西湖闻名，是著名的旅游城市。",
            "古月方源的弟弟叫古月方正",
            "古月方源号称大爱仙尊, 但是被天庭污蔑为练天魔尊",
            "古月方源的忠实迷弟是风天语",
        ]
        documents = [
            *_documents,
            "Project Initialization: Always start with a version control system like Git. Create a clear directory structure, e.g., 'src' for source code, 'tests' for tests, and 'docs' for documentation. Use a virtual environment to manage dependencies.",
            "PDF Processing (Python): For PDF processing in Python, popular libraries include PyPDF2 for text extraction and metadata, and pdf2image for converting pages into images for OCR.",
            "API Development (Python): When building an API with Python, FastAPI is a modern, high-performance choice. It includes automatic data validation and documentation.",
            "OCR (Optical Character Recognition): Tesseract is a powerful open-source OCR engine. For cloud-based solutions, Google Cloud Vision and AWS Textract offer high-accuracy OCR APIs.",
            "Markdown Conversion: The 'markdown' library in Python can convert markdown text to HTML. To create a markdown file, you can simply write to a file with a '.md' extension.",
            "Command-Line Interface (CLI): Python's `argparse` module is the standard way to create a user-friendly command-line interface. For more complex CLIs, libraries like 'Click' or 'Typer' are excellent choices.",
            "Testing (Python): Pytest is the recommended framework for testing Python code. It supports simple unit tests and complex functional testing, and has a rich ecosystem of plugins.",
            "Dependencies Management: A `requirements.txt` file is the standard way to list project dependencies. For more robust dependency management, consider using tools like Poetry or Pipenv."
        ]
        self.documents = documents
        for i, doc in enumerate(documents):
            doc = self.get_cut_words(doc)
            response = ollama.embed(model=self.embedding_model, input=doc)
            # The key is 'embedding', not 'embeddings'
            embeddings = response.get("embeddings")
            if embeddings:
                self.collection.add(
                    ids=[str(i)],
                    embeddings=embeddings,
                    documents=[doc]
                )

    def retrieve(self, query: str, n_results: int = 3) -> str:
        """
        Retrieves relevant information. Uses jieba to segment Chinese queries.
        """
        query = self.get_cut_words(query, debug=True)
        print("~~~~~ query:", query)
        response = ollama.embed(model=self.embedding_model, input=query)
        embeddings = response.get("embeddings")
        if not embeddings:
            return "*** Could not generate embedding for the query."

        results = self.collection.query(
            query_embeddings=embeddings,
            n_results=n_results
        )
        
        # retrieved_docs = results.get('documents', [[]])[0]
        ids = results.get('ids', [[]])[0]
        if not ids:
            return "No relevant technical information found."

        # docs = self.collection.get(ids=ids)
        # retrieved_docs = [doc for doc in self.collection.get(ids=ids)]
        # retrieved_docs = self.documents[ids]
        retrieved_docs = [self.documents[int(i)] for i in ids]
        return "\n".join(retrieved_docs)

if __name__ == '__main__':
    # Example usage:
    print("--- Initializing PlanningRAG...")
    rag = PlanningRAG()
    
    print("\n--- Testing PDF Query ---")
    # pdf_query = "How should I handle PDF files in my project?"
    # pdf_query = "介绍下古月方源"
    pdf_query = "介绍一下方源"
    # pdf_query = "风天语是谁"
    tech_info_pdf = rag.retrieve(pdf_query)
    print(f"--- Query: {pdf_query}")
    print("--- Retrieved Information:")
    print(tech_info_pdf)

    exit()

    # print("\n--- Testing API Query ---")
    # api_query = "What's a good way to build an API in Python?"
    # tech_info_api = rag.retrieve(api_query)
    # print(f"--- Query: {api_query}")
    # print("--- Retrieved Information:")
    # print(tech_info_api)

